#include "PCGSolver.h"
#include "mkl_rci.h"
#include "mkl_blas.h"
#include "mkl_spblas.h"
#include "mkl_service.h"
#include "DecoLogger.h"
#include <list>
#include <vector>
#include <math.h>

PCGSolver::PCGSolver() : mPrecon(NULL)
{
	//mkl_set_num_threads( 1 );
	//mkl_domain_set_num_threads( 16, MKL_ALL );
	//mkl_domain_set_num_threads( 8, MKL_BLAS );
}
PCGSolver::~PCGSolver()
{
	if (mPrecon)
		delete[] mPrecon;
}

void PCGSolver::Solve(double* result, bool usePreconditioner)
{
	if (!mMatrix->IsSymmetric())
	{
		(*DecoLogger::GetSingleton()) << "PCG solver can only solve symmetric matrix!\n";
	}
	if (usePreconditioner)
		calculatePreconditioner();
	char tr='l';
	char matdes[4];
	matdes[0] = 'd';
	matdes[1] = 'l';
	matdes[2] = 'n';
	matdes[3] = 'f';
	int dim = mMatrix->n;
	MKL_INT rci_request, itercount;
	MKL_INT ipar[128];
	double dpar[128];
	double* a = &(mMatrix->values[0]);
	int* ia = &(mMatrix->rowIndex[0]);
	int* ja = &(mMatrix->columns[0]);


	double *tmp = new double[4 * dim];
	double *trvec = new double[dim];
	double one = 1.0;

	/*---------------------------------------------------------------------------*/
	/* Initialize the initial guess                                              */
	/*---------------------------------------------------------------------------*/
	for(int i = 0; i < dim;i++) 
		result[i] = 0;
	/*---------------------------------------------------------------------------*/
	/* Initialize the solver                                                     */
	/*---------------------------------------------------------------------------*/
	dcg_init(&dim, result, mRHS, &rci_request, ipar, dpar, tmp);
	if (rci_request != 0) goto failure;
	/*---------------------------------------------------------------------------*/
	/* Set the desired parameters:                                               */
	/* LOGICAL parameters:                                                       */
	/* do residual stopping test                                                 */
	/* do not request for the user defined stopping test                         */
	/* DOUBLE parameters                                                         */
	/*---------------------------------------------------------------------------*/
	ipar[1] = 6;
	ipar[4] = 1000;
	ipar[5] = 1;
	ipar[6] = 1;
	ipar[8] = 1;
	ipar[9] = 0;
	ipar[10] = usePreconditioner? 1 : 0;

	dpar[0] = mErrorThreshold;

	/*---------------------------------------------------------------------------*/
	/* Check the correctness and consistency of the newly set parameters         */
	/*---------------------------------------------------------------------------*/
	dcg_check(&dim, result, mRHS, &rci_request, ipar, dpar, tmp);
	if (rci_request!=0) goto failure;
	/*---------------------------------------------------------------------------*/
	/* Compute the solution by RCI (P)CG solver without preconditioning          */
	/* Reverse Communications starts here                                        */
	/*---------------------------------------------------------------------------*/
rci: dcg(&dim, result, mRHS, &rci_request, ipar, dpar, tmp);
	/*---------------------------------------------------------------------------*/
	/* If rci_request=0, then the solution was found with the required precision */
	/*---------------------------------------------------------------------------*/
	if (rci_request==0) goto getsln;
	/*---------------------------------------------------------------------------*/
	/* If rci_request=1, then compute the vector A*tmp[0]                 */
	/* and put the result in vector tmp[n]                                       */
	/*---------------------------------------------------------------------------*/

	if (rci_request==1)
	{
		//tr = 'n';
		//mkl_dcsrgemv(&tr, &dim, a, ia, ja, tmp, &tmp[dim]);
		mkl_dcsrsymv(&tr, &dim, a, ia, ja, tmp, &tmp[dim]);
		//mkl_cspblas_dcsrsymv(&tr, &dim, a, ia, ja, tmp, &tmp[dim]);

		//for (int i = 0; i < dim; ++i)
		//{
		//	(*DecoLogger::GetSingleton()) << i << ":  " << tmp[i] << "  " << tmp[dim + i] << "\n";
		//}
		goto rci;
	}

	/*---------------------------------------------------------------------------*/
	/* If rci_request=3, then compute apply the preconditioner matrix C_inverse  */
	/* on vector TMP(:,3) and put the result in vector tmp[3*n]                  */
	/*---------------------------------------------------------------------------*/

	if (rci_request == 3)
	{
		//mkl_dcsrsv(&matdes[2], &dim, &one, matdes, a, ja, ia, &ia[1], &tmp[2*dim], &tmp[3*dim]);

		//goto rci;

		char cvar1='L';
		char cvar='T';
		char cvar2='N';
		mkl_dcsrtrsv(&cvar1,&cvar,&cvar2,&dim,mPrecon,ia,ja,&tmp[2*dim],trvec);
		cvar1='L';
		cvar='N';
		cvar2='N';
		mkl_dcsrtrsv(&cvar1,&cvar,&cvar2,&dim,mPrecon,ia,ja,trvec,&tmp[3*dim]);
		goto rci;
	}
	/*---------------------------------------------------------------------------*/
	/* If rci_request=anything else, then dcg subroutine failed                  */
	/* to compute the solution vector: solution[n]                               */
	/*---------------------------------------------------------------------------*/
	goto failure;
	/*---------------------------------------------------------------------------*/
	/* Reverse Communication ends here                                           */
	/* Get the current iteration number into itercount                           */
	/*---------------------------------------------------------------------------*/
getsln: dcg_get(&dim, result, mRHS, &rci_request, ipar, dpar, tmp, &itercount);

	(*DecoLogger::GetSingleton()) << itercount << " number of iterations used to solve " << dim << " unknowns.\n";
	delete[] tmp;
	delete[] trvec;

	return;

failure: (*DecoLogger::GetSingleton()) << "This example FAILED with " << ipar[3] << " iterations and the solver has returned the ERROR code " << rci_request << "\n";
	delete[] tmp;
	delete[] trvec;

	return;
}

void PCGSolver::calculatePreconditioner() 
{
	if (mPrecon)
		delete[] mPrecon;
	
	int numEntries = mMatrix->GetNumNonZeroEntries();
	mPrecon = new double[numEntries];	
	memset(mPrecon, 0, numEntries * sizeof(double));

	double* a = &(mMatrix->values[0]);
	int* ia = &(mMatrix->rowIndex[0]);
	int* ja = &(mMatrix->columns[0]);

	int dim = mMatrix->n;

	for (int ithRow = 1; ithRow <= dim; ++ithRow)
	{
		int ithRowStartIdx = ia[ithRow - 1] - 1;
		int ithRowEndIdx = ia[ithRow] - 1;

		for (int idx = ithRowStartIdx; idx < ithRowEndIdx; ++idx)
		{
			int jthCol = ja[idx];
			mPrecon[idx] = a[idx];
			if (ithRow == jthCol)
			{
				for (int k = ithRowStartIdx; k < idx; ++k)
				{
					mPrecon[idx] -= mPrecon[k] * mPrecon[k];
				}
				mPrecon[idx] = sqrt(mPrecon[idx]);
			}
			else
			{
				int jthRowStartIdx = ia[jthCol - 1] - 1;
				int jthRowEndIdx = ia[jthCol] - 1;
				
				for (int k1 = ithRowStartIdx; k1 < ithRowEndIdx; ++k1)
				{
					for (int k2 = jthRowStartIdx; k2 < jthRowEndIdx; ++k2)
					{
						if (ja[k1] == ja[k2] && ja[k1] < jthCol)
						{
							mPrecon[idx] -= mPrecon[k1] * mPrecon[k2];
						}
					}
				}
				double Ljj = 1.0;
				for (int k2 = jthRowStartIdx; k2 < jthRowEndIdx; ++k2)
				{
					if (ja[k2] == jthCol)
						Ljj = mPrecon[k2];
				}
				mPrecon[idx] /= Ljj;
			}
		}
	}

	//std::list<int> notFinishedRows;
	//std::vector<std::list<int>::iterator> finishedRowIt;

	//for (int ithRow = 1; ithRow <= dim; ++ithRow)
	//{
	//	notFinishedRows.push_back(ithRow);
	//}
	//for (int ithColEntry = 0; ithColEntry < dim; ++ithColEntry)
	//{
	//	if (notFinishedRows.empty())
	//		break;
	//	finishedRowIt.clear();
	//	for (std::list<int>::iterator rowIt = notFinishedRows.begin(); rowIt != notFinishedRows.end(); ++rowIt)
	//	{
	//		int ithRow = *rowIt;
	//		int rowStartIdx = ia[ithRow - 1] - 1; 
	//		int rowEndIdx = ia[ithRow] - 1;
	//		int numColEntries = rowEndIdx - rowStartIdx;
	//		if (ithColEntry == numColEntries - 1)
	//		{
	//			finishedRowIt.push_back(rowIt);
	//		}
	//		int idx = rowStartIdx + ithColEntry;
	//		int jthCol = ja[idx];
	//		mPrecon[idx] = a[idx];
	//		if (ithRow == jthCol)
	//		{
	//			for (int k = rowStartIdx; k < idx; ++k)
	//			{
	//				mPrecon[idx] -= mPrecon[k] * mPrecon[k];
	//			}
	//			mPrecon[idx] = sqrt(mPrecon[idx]);
	//		}
	//		else
	//		{
	//			int ithRowStartIdx = ia[ithRow - 1] - 1;
	//			int jthRowStartIdx = ia[jthCol - 1] - 1;
	//			//for (int k = 0; k )
	//		}

	//	}
	//	int numFinishedRows = static_cast<int>(finishedRowIt.size());
	//	for (int i = 0; i < numFinishedRows; ++i)
	//	{
	//		notFinishedRows.erase(finishedRowIt[i]);
	//	}
	//}



}